Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly upgrades the hardware parsing and provisioning capabilities within keras-remote. It allows users to specify hardware requirements more abstractly, such as "gpu-4" or "tpu-512", and intelligently selects the most suitable accelerator, prioritizing newer generations. This change streamlines the developer experience by abstracting away complex GKE hardware topologies while maintaining backward compatibility and expanding the range of supported hardware configurations.

Highlights

• Dynamic Hardware Provisioning: Introduced logic to automatically provision GPUs and TPUs based on generic requests like "gpu-16" or "tpu-512", removing the need for exact GKE topology knowledge.
• Generation-Aware Prioritization: Implemented a preference-based fallback mechanism for TPUs and GPUs, ensuring that newer hardware generations are prioritized when generic requests are made.
• Canonical Accelerator Aliases: Added support for aliases like "v5e" and "ghostlite" to map to canonical TPU names, ensuring consistent backend labeling.
• Expanded Hardware Support: Updated GPU and TPU specifications to include new hardware types (P4, P100) and expanded available counts/topologies for existing ones (L4, A100, V2, V3, V4, V5litepod).
• Improved Parsing Robustness: Refined regex patterns for accelerator parsing to prevent false positives and enhance efficiency, along with comprehensive unit test updates.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

/gemini review

/gemini review

Just a note, we presently create all GPU node pools with count=1. Would leave it to be your call if you'd like to fix it as part of this change.

Presently, GpuSpec stores a single machine_type for all counts. We need to follow the TpuSpec model which maps each count to its own machine type.

We only need to map the int count to an str machine_type (instead of requiring a GpuTopology class), since only machine_type varies per count.

Something like:

@dataclass(frozen=True)
class GpuSpec:
  gke_label: str
  counts: dict[int, str]    # count -> machine_type

Then the GPUS registry will need to be updated as:

GPUS = {
  "a100": GpuSpec("nvidia-tesla-a100", {
    1: "a2-highgpu-1g",
    2: "a2-highgpu-2g",
    # ...
  }),
  "h100": GpuSpec("nvidia-h100-80gb", {
    1: "a3-highgpu-1g",
    # ...
  }),
# ...
}

Then we can update the guest_accelerators count in node pool creation:

guest_accelerators=[
  gcp.container.NodePoolNodeConfigGuestAcceleratorArgs(
    type=gpu.gke_label,
    count=gpu.count,
  ),
],

remote/keras_remote/cli/infra/program.py

Lines 233 to 236 in 70bd83e

	gcp.container.NodePoolNodeConfigGuestAcceleratorArgs(
	type=gpu.gke_label,
	count=1,
	),

Just a note, we presently create all GPU node pools with count=1. Would leave it to be your call if you'd like to fix it as part of this change.

Presently, GpuSpec stores a single machine_type for all counts. We need to follow the TpuSpec model which maps each count to its own machine type.

We only need to map the int count to an str machine_type (instead of requiring a GpuTopology class), since only machine_type varies per count.

Something like:

@dataclass(frozen=True)
class GpuSpec:
  gke_label: str
  counts: dict[int, str]    # count -> machine_type

Then the GPUS registry will need to be updated as:

GPUS = {
  "a100": GpuSpec("nvidia-tesla-a100", {
    1: "a2-highgpu-1g",
    2: "a2-highgpu-2g",
    # ...
  }),
  "h100": GpuSpec("nvidia-h100-80gb", {
    1: "a3-highgpu-1g",
    # ...
  }),
# ...
}

Then we can update the guest_accelerators count in node pool creation:

guest_accelerators=[
  gcp.container.NodePoolNodeConfigGuestAcceleratorArgs(
    type=gpu.gke_label,
    count=gpu.count,
  ),
],

Good catch! I have refactored GpuSpecto mirror TpuSpec's behavior. It now uses a counts: dict[int, str] dictionary to map explicit integer counts directly to their corresponding optimal GKE machine type (e.g. mapping l4 count 8 perfectly to g2-standard-96). I updated keras_remote/cli/infra/program.py so guest_accelerators dynamically injects count=gpu.count instead of hardcoding 1.

FYI: The local e2e tests are passing.